// Copyright (c) Lawrence Livermore National Security, LLC and other VisIt
// Project developers.  See the top-level LICENSE file for dates and other
// details.  No copyright assignment is required to contribute to VisIt.

// ************************************************************************* //
//                            avtALSFileFormat.C                           //
// ************************************************************************* //

#include <avtALSFileFormat.h>

#include <string>

#include <vtkFloatArray.h>
#include <vtkRectilinearGrid.h>
#include <vtkStructuredGrid.h>
#include <vtkUnstructuredGrid.h>

#include <avtDatabaseMetaData.h>

#include <Expression.h>

#include <InvalidVariableException.h>

#include <hdf5.h>
#include <visit-hdf5.h>
#include <avtParallel.h>
#include <vtkDoubleArray.h>

using     std::string;


/// iterate over groups.
herr_t group_info(hid_t loc_id, const char *name, void *opdata)
{
    avtALSFileFormat* format = (avtALSFileFormat*)opdata;

    H5G_stat_t statbuf;

    /*
     * Get type of the object and display its name and type.
     * The name of the object is passed to this function by
     * the Library. Some magic :-)
     */
    H5Gget_objinfo(loc_id, name, 0, &statbuf);
    switch (statbuf.type) {
    case H5G_GROUP:
         format->SetGroupName(name);
         break;
    default:
         printf(" Unable to identify an object ");
    }
    return 0;
 }



// ****************************************************************************
//  Method: avtALSFileFormat constructor
//
//  Programmer: hari -- generated by xml2avt
//  Creation:   Mon Apr 14 16:36:42 PST 2014
//
// ****************************************************************************

avtALSFileFormat::avtALSFileFormat(const char *filename)
    : avtSTMDFileFormat(&filename, 1)
{
    m_dataType = STANDARD; //STANDARD, TOMO, OTHER..
    m_filename = filename;
    m_width = m_height = m_slices = 0;
    m_initialized = false;
    m_ver = 0;
}

void
avtALSFileFormat::Initialize() {

    if(m_initialized) {
        return;
    }

    if(PAR_Size() == 1) {
        InitializeHeader();
    }
    else {

        int rank = PAR_Rank();
        if(PAR_UIProcess()) {

            bool success = InitializeHeader();

            BroadcastBool(success);

            if(!success) {
                return;
            }

            intVector dims;
            dims.push_back(m_width);
            dims.push_back(m_height);
            dims.push_back(m_slices);

            BroadcastIntVector(dims, rank);
            int dataType = (int)m_dataType;
            BroadcastInt(dataType);

            BroadcastString(m_groupName, rank);

            if(m_dataType == STANDARD) {
                BroadcastStringVector(m_datasetNames, rank);
            }

        } else {
            bool success = false;

            BroadcastBool(success);

            if(!success) {
                return;
            }

            m_initialized = true;

            intVector dims;
            BroadcastIntVector(dims, rank);

            m_width = dims[0];
            m_height = dims[1];
            m_slices = dims[2];
            int dataType = -1;
            BroadcastInt(dataType);
            m_dataType = (ALSDataType)dataType;
            BroadcastString(m_groupName, rank);

            if(m_dataType == STANDARD) {
                BroadcastStringVector(m_datasetNames, rank);
            }
        }
    }

    //std::cout << PAR_Rank() << " " << m_groupName << " " << m_width << " " << m_height << " " << m_slices
    //          << " " << m_datasetNames.size() << std::endl;
}


bool
avtALSFileFormat::InitializeTomoHeader(hid_t file) {

    /// get exchange id
    std::string four_d_data_str = "/exchange/data/4Ddata";
    std::string fraction_data_str = "/exchange/data/fraction";
    bool four_d_data = false, fraction_data = false;
   
    four_d_data = H5Lexists(file, four_d_data_str.c_str(), H5P_DEFAULT);

    if(!four_d_data) { 
         fraction_data = H5Lexists(file, fraction_data_str.c_str(), H5P_DEFAULT);
    }

    if(!four_d_data && !fraction_data) {
        return false;
    }

    m_groupName = fraction_data ? fraction_data_str : four_d_data_str;

    hid_t dataset_id = H5Dopen(file, m_groupName.c_str(), H5P_DEFAULT);
    
    hsize_t dims[4] = {0,0,0,0};
    hid_t filespace = H5Dget_space(dataset_id);
    hid_t rank      = H5Sget_simple_extent_ndims(filespace);
    hid_t status_n  = H5Sget_simple_extent_dims(filespace, dims, NULL);

    std::cout << dims[0] << " " << dims[1] << " " << dims[2] << " " << dims[3] << std::endl;

    m_dataType = TOMO;
    m_width = dims[1];
    m_height = dims[2];
    m_slices = dims[3];
    m_ver = 0;
    m_initialized = true;

    H5Dclose(dataset_id);

    return true;
}

bool
avtALSFileFormat::InitializeStandardHeader(hid_t file) {
    hid_t status;

    /// else STANDARD
    m_dataType = STANDARD;
    hid_t group_id = H5Gopen1(file, m_groupName.c_str());

    hsize_t num_objs = 0;
    H5Gget_num_objs(group_id, &num_objs);
    /// get nslices attributes to read in the number of slices.

    int dim1, dim2, dim3;
    /// read only the first one, perform a sanity check against all
    /// if really necessary.

    for(size_t i = 0; i < num_objs; ++i)
    {
        char name[1024];
        H5Gget_objname_by_idx(group_id, i, name, 1024);

        hid_t attr_id;

        hid_t dataset_id = H5Dopen1(group_id, name);

        int sdim1, sdim2, sdim3;

        attr_id = H5Aopen(dataset_id, "dim1", H5P_DEFAULT);
        H5Aread(attr_id, H5T_NATIVE_INT, &sdim1);
        H5Aclose(attr_id);

        attr_id = H5Aopen(dataset_id, "dim2", H5P_DEFAULT);
        H5Aread(attr_id, H5T_NATIVE_INT, &sdim2);
        H5Aclose(attr_id);

        attr_id = H5Aopen(dataset_id, "dim3", H5P_DEFAULT);
        H5Aread(attr_id, H5T_NATIVE_INT, &sdim3);
        H5Aclose(attr_id);

        H5Dclose(dataset_id);

        std::string sname = name;

        /// all dataset names should match
        if(sname.find(m_groupName) == std::string::npos) {
            //debug4() << "Group Name does not match Dataset Name" << std::endl;
            std::cerr << "Group Name does not match Dataset Name: Expected Part " << m_groupName
                      << " Got " << sname << " " << std::endl;
            status = H5Fclose(file);
            return false;
        }

        if(i == 0) {
            dim1 = sdim1;
            dim2 = sdim2;
            dim3 = sdim3;
        } else {
            /// if dimensions don't match
            if(dim1 != sdim1 || dim2 != sdim2 || dim3 != sdim3) {
                //debug4() << "Not ALS File Dimensions don't match" << std::endl;
                std::cerr << "Not ALS File Dimensions don't match" << std::endl;
                status = H5Fclose(file);
                return false;
            }
        }

        //std::cout << name << std::endl;
        //std::cout << dim1 << " " << dim2 << " " << dim3 << std::endl;
        m_datasetNames.push_back(sname);
    }

    ///shift by 1 dimension to make last dimension total number of slices..
    m_width = dim2;
    m_height = dim3;
    m_slices = num_objs;

    //std::cout << dim2 << " " << dim3 << " " << num_objs << std::endl;

    H5Gclose(group_id);

    m_initialized = true;
    return true;
}

bool
avtALSFileFormat::InitializeHeader() {
    hid_t           file;           /* Handle */
    herr_t          status;

    (void) status; /// remove warning..

    file = H5Fopen (m_filename.c_str(), H5F_ACC_RDONLY, H5P_DEFAULT);

    /// main group seems to be based on filename..

    H5Giterate(file, "/", NULL, group_info, this);

    /// if group is not set..
    if(m_groupName.length() == 0) {
        std::cerr << "Group Name not found, must not be ALS dataset" << std::endl;
        status = H5Fclose(file);
        return false;
    }

    bool result = false;

    if (m_groupName == "exchange") {
        result = InitializeTomoHeader(file);
    } else {
        result = InitializeStandardHeader(file);
    }

    H5Fclose(file);
    return result;
}

vtkDataArray*
avtALSFileFormat::GetTomoDataSet(hid_t file) {
    vtkDoubleArray* array = vtkDoubleArray::New();

    std::cout << m_groupName << " " << m_width << " " << m_height << " " << m_slices << std::endl;

    hsize_t     dims_out[4];           /* dataset dimensions */
    hid_t dataset_id = H5Dopen(file, m_groupName.c_str(), H5P_DEFAULT);
    hid_t dataspace = H5Dget_space(dataset_id);
    hid_t rank      = H5Sget_simple_extent_ndims(dataspace);
    hid_t status_n  = H5Sget_simple_extent_dims(dataspace, dims_out, NULL);

    (void) rank;
    (void) status_n;

    hid_t       memspace;
    herr_t      status;

    hsize_t     count[4];              /* size of the hyperslab in the file */
    hsize_t     offset[4];             /* hyperslab offset in the file */

    hsize_t     dimsm[3];              /* memory space dimensions */
    hsize_t     count_out[3];          /* size of the hyperslab in memory */
    hsize_t     offset_out[3];         /* hyperslab offset in memory */

    offset[0] = 0; /// pick offset dataset..
    offset[1] = 0;
    offset[2] = 0;
    offset[3] = 0;

    count[0]  = 1;
    count[1]  = dims_out[1];
    count[2]  = dims_out[2];
    count[3]  = dims_out[3];

    status = H5Sselect_hyperslab (dataspace, H5S_SELECT_SET, offset, NULL,
                                  count, NULL);

    if(PAR_Size() == 1) {
        /*
         * Define the memory dataspace.
         */
        dimsm[0] = dims_out[1];
        dimsm[1] = dims_out[2];
        dimsm[2] = dims_out[3];

        memspace = H5Screate_simple (3, dimsm, NULL);

        /*
         * Define memory hyperslab.
         */
        offset_out[0] = 0;
        offset_out[1] = 0;
        offset_out[2] = 0;

        count_out[0]  = dims_out[1];
        count_out[1]  = dims_out[2];
        count_out[2]  = dims_out[3];

        status = H5Sselect_hyperslab (memspace, H5S_SELECT_SET, offset_out, NULL,
                                      count_out, NULL);

        array->SetNumberOfTuples(m_width*m_height*m_slices);
        H5Dread(dataset_id, H5T_NATIVE_DOUBLE, memspace, dataspace, H5P_DEFAULT, array->GetVoidPointer(0));
    } else {
        int m_offset = m_slices / PAR_Size();

        int m_start = PAR_Rank()*m_offset;
        int m_end = m_start + m_offset;

        if(PAR_Rank() == PAR_Size()-1) {
            m_end = m_slices;
        }

        int zslices = m_end-m_start + 1;

        array->SetNumberOfTuples(m_width * m_height * zslices);

        dimsm[0] = 0;
        dimsm[1] = dims_out[1];
        dimsm[2] = dims_out[2];
        //dimsm[3] = zslices;

        memspace = H5Screate_simple (4, dimsm, NULL);

        /*
         * Define memory hyperslab.
         */
        offset_out[0] = 0;
        offset_out[1] = 0;
        offset_out[2] = 0;
        //offset_out[3] = 0;

        count_out[0]  = 0;
        count_out[1]  = dims_out[1];
        count_out[2]  = dims_out[2];
        //count_out[3]  = zslices;

        status = H5Sselect_hyperslab (memspace, H5S_SELECT_SET, offset_out, NULL,
                                      count_out, NULL);

        H5Dread(dataset_id, H5T_NATIVE_DOUBLE, memspace, dataspace, H5P_DEFAULT, array->GetVoidPointer(0));
    }

    H5Dclose(dataset_id);
    status = H5Fclose(file);

    return array;
}

vtkDataArray*
avtALSFileFormat::GetStandardDataSet(hid_t file) {
    vtkFloatArray* array = vtkFloatArray::New();

    /// default STANDARD
    hid_t group_id = H5Gopen(file, m_groupName.c_str(), H5P_DEFAULT);


    if(PAR_Size() == 1) {
        array->SetNumberOfTuples(m_width * m_height * m_slices);
        for(size_t i = 0; i < m_datasetNames.size(); ++i)
        {
            hid_t dataset_id = H5Dopen(group_id, m_datasetNames[i].c_str(), H5P_DEFAULT);

            ///dim1 should be 1, dim2 and dim3 is the size of the dataset
            H5Dread(dataset_id, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT,array->GetVoidPointer(i*m_width*m_height));
            H5Dclose(dataset_id);
        }
    } else {
        int m_offset = m_slices / PAR_Size();

        int m_start = PAR_Rank()*m_offset;
        int m_end = m_start + m_offset + 1;

        if(PAR_Rank() == PAR_Size()-1) {
            m_end = m_slices;
        }

        int zslices = m_end-m_start;

        array->SetNumberOfTuples(m_width * m_height * zslices);

        for(size_t i = m_start, j = 0; i < m_end; ++i, ++j)
        {
            hid_t dataset_id = H5Dopen(group_id, m_datasetNames[i].c_str(), H5P_DEFAULT);

            ///dim1 should be 1, dim2 and dim3 is the size of the dataset
            H5Dread(dataset_id, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT,array->GetVoidPointer(j*m_width*m_height));
            H5Dclose(dataset_id);
        }
    }

    H5Gclose(group_id);

    return array;
}

vtkDataArray*
avtALSFileFormat::GetDataSet() {

    if(!m_initialized) {
        return NULL;
    }

    hid_t           file;           /* Handle */
    herr_t          status;

    (void) status; /// remove warning..


    file = H5Fopen (m_filename.c_str(), H5F_ACC_RDONLY, H5P_DEFAULT);

    vtkDataArray* array = NULL;

    if(m_dataType == TOMO) {
        array = GetTomoDataSet(file);
    } else {
        array = GetStandardDataSet(file);
    }

    status = H5Fclose(file);

    return array;
}

// ****************************************************************************
//  Method: avtALSFileFormat::FreeUpResources
//
//  Purpose:
//      When VisIt is done focusing on a particular timestep, it asks that
//      timestep to free up any resources (memory, file descriptors) that
//      it has associated with it.  This method is the mechanism for doing
//      that.
//
//  Programmer: hari -- generated by xml2avt
//  Creation:   Mon Apr 14 16:36:42 PST 2014
//
// ****************************************************************************

void
avtALSFileFormat::FreeUpResources(void)
{
}


// ****************************************************************************
//  Method: avtALSFileFormat::PopulateDatabaseMetaData
//
//  Purpose:
//      This database meta-data object is like a table of contents for the
//      file.  By populating it, you are telling the rest of VisIt what
//      information it can request from you.
//
//  Programmer: hari -- generated by xml2avt
//  Creation:   Mon Apr 14 16:36:42 PST 2014
//
// ****************************************************************************

void
avtALSFileFormat::PopulateDatabaseMetaData(avtDatabaseMetaData *md)
{
    /// both MDServer & Engine call this?
    Initialize();

    if(!m_initialized) {
        return;
    }

    std::string meshname = "mesh";

    avtMeshType mt = AVT_RECTILINEAR_MESH;

    int dimensions = 3;
    int nblocks = 1; //possibly change this for parallel implementation?
    int block_origin = 0;
    int spatial_dimension = dimensions;
    int topological_dimension = dimensions;
    const double *extents = NULL;

    //const int *bounds = NULL;
    AddMeshToMetaData(md, meshname, mt, extents, nblocks, block_origin,
                      spatial_dimension, topological_dimension);


    //Add Scalar information
    std::string mesh_for_this_var = meshname; // ??? -- could be multiple meshes
    avtCentering cent = AVT_NODECENT;

    /// TODO: CHECK IF THIS ALWAYS HAS 3..
    if(m_dataType == TOMO) {
        AddScalarVarToMetaData(md, "intensity1", mesh_for_this_var, cent);
        AddScalarVarToMetaData(md, "intensity2", mesh_for_this_var, cent);
        AddScalarVarToMetaData(md, "intensity3", mesh_for_this_var, cent);
    } else {

        std::string varname = "intensity";
        AddScalarVarToMetaData(md, varname, mesh_for_this_var, cent);
    }
    //telling the meta data that we can do dynamic domain decomposition
    md->SetFormatCanDoDomainDecomposition(true);
}


// ****************************************************************************
//  Method: avtALSFileFormat::GetMesh
//
//  Purpose:
//      Gets the mesh associated with this file.  The mesh is returned as a
//      derived type of vtkDataSet (ie vtkRectilinearGrid, vtkStructuredGrid,
//      vtkUnstructuredGrid, etc).
//
//  Arguments:
//      domain      The index of the domain.  If there are NDomains, this
//                  value is guaranteed to be between 0 and NDomains-1,
//                  regardless of block origin.
//      meshname    The name of the mesh of interest.  This can be ignored if
//                  there is only one mesh.
//
//  Programmer: hari -- generated by xml2avt
//  Creation:   Mon Apr 14 16:36:42 PST 2014
//
// ****************************************************************************

vtkDataSet *
avtALSFileFormat::GetMesh(int domain, const char *meshname)
{
    if(!m_initialized) {
        return NULL;
    }

    vtkRectilinearGrid* dataset = vtkRectilinearGrid::New();

    vtkFloatArray* xArray = vtkFloatArray::New();
    xArray->SetNumberOfTuples(m_width);
    xArray->SetNumberOfComponents(1);

    for(size_t i = 0; i < m_width; ++i)
        xArray->SetTuple1(i, i);

    vtkFloatArray* yArray = vtkFloatArray::New();
    yArray->SetNumberOfTuples(m_height);
    yArray->SetNumberOfComponents(1);

    for(size_t i = 0; i < m_height; ++i)
        yArray->SetTuple1(i, i);

    vtkFloatArray* zArray = vtkFloatArray::New();
    int zslices = m_slices;

    if(PAR_Size() == 1) {
        zArray->SetNumberOfTuples(zslices);
        zArray->SetNumberOfComponents(1);

        for(size_t i = 0; i < zslices; ++i)
            zArray->SetTuple1(i, i);
    } else {

        int m_offset = m_slices / PAR_Size();

        int m_start = PAR_Rank()*m_offset;
        int m_end = m_start + m_offset + 1;

        if(PAR_Rank() == PAR_Size()-1) {
            m_end = m_slices;
        }

        zslices = m_end-m_start;

        zArray->SetNumberOfTuples(zslices);
        zArray->SetNumberOfComponents(1);

        //std::cout << PAR_Rank() << " " << m_start << " " << m_end << " " << m_slices << std::endl;

        for(size_t i = m_start; i < m_end; ++i)
            zArray->SetTuple1(i-m_start, i);
    }

    dataset->SetXCoordinates(xArray);
    dataset->SetYCoordinates(yArray);
    dataset->SetZCoordinates(zArray);

    dataset->SetDimensions(m_width, m_height, zslices);

    xArray->Delete();
    yArray->Delete();
    zArray->Delete();

    return dataset;
}


// ****************************************************************************
//  Method: avtALSFileFormat::GetVar
//
//  Purpose:
//      Gets a scalar variable associated with this file.  Although VTK has
//      support for many different types, the best bet is vtkFloatArray, since
//      that is supported everywhere through VisIt.
//
//  Arguments:
//      domain     The index of the domain.  If there are NDomains, this
//                 value is guaranteed to be between 0 and NDomains-1,
//                 regardless of block origin.
//      varname    The name of the variable requested.
//
//  Programmer: hari -- generated by xml2avt
//  Creation:   Mon Apr 14 16:36:42 PST 2014
//
// ****************************************************************************

vtkDataArray *
avtALSFileFormat::GetVar(int domain, const char *varname)
{
     if(m_dataType == TOMO) {
         std::string name = varname;
         if(name == "intensity1") m_ver = 0;
         if(name == "intensity2") m_ver = 1;
         if(name == "intensity3") m_ver = 2;

     }
     return GetDataSet();
}


// ****************************************************************************
//  Method: avtALSFileFormat::GetVectorVar
//
//  Purpose:
//      Gets a vector variable associated with this file.  Although VTK has
//      support for many different types, the best bet is vtkFloatArray, since
//      that is supported everywhere through VisIt.
//
//  Arguments:
//      domain     The index of the domain.  If there are NDomains, this
//                 value is guaranteed to be between 0 and NDomains-1,
//                 regardless of block origin.
//      varname    The name of the variable requested.
//
//  Programmer: hari -- generated by xml2avt
//  Creation:   Mon Apr 14 16:36:42 PST 2014
//
// ****************************************************************************

vtkDataArray *
avtALSFileFormat::GetVectorVar(int domain, const char *varname)
{
    return NULL;
}
